#!/usr/bin/env python3
import argparse
import contextlib
import io
import json
import signal
import socket
import os
import subprocess
import sys
import tempfile
import time
import unittest
from pathlib import Path


DEFAULT_VENV = "~/.local/share/tsrs-kokoro/venv"
DEFAULT_VOICE_ID = "af_heart"
DEFAULT_SERVER_IDLE_TIMEOUT_SECONDS = 1800
SERVER_PROTOCOL_VERSION = 1
VOICE_IDS = [
    "af_alloy",
    "af_aoede",
    "af_bella",
    "af_heart",
    "af_jessica",
    "af_kore",
    "af_nicole",
    "af_nova",
    "af_river",
    "af_sarah",
    "af_sky",
    "am_adam",
    "am_echo",
    "am_eric",
    "am_fenrir",
    "am_liam",
    "am_michael",
    "am_onyx",
    "am_puck",
    "am_santa",
    "bf_alice",
    "bf_emma",
    "bf_isabella",
    "bf_lily",
    "bm_daniel",
    "bm_fable",
    "bm_george",
    "bm_lewis",
    "ef_dora",
    "em_alex",
    "em_santa",
    "ff_siwis",
    "hf_alpha",
    "hf_beta",
    "hm_omega",
    "hm_psi",
    "if_sara",
    "im_nicola",
    "jf_alpha",
    "jf_gongitsune",
    "jf_nezumi",
    "jf_tebukuro",
    "jm_kumo",
    "pf_dora",
    "pm_alex",
    "pm_santa",
    "zf_xiaobei",
    "zf_xiaoni",
    "zf_xiaoxiao",
    "zf_xiaoyi",
    "zm_yunjian",
    "zm_yunxi",
    "zm_yunxia",
    "zm_yunyang",
]


def expanded_path(value):
    return Path(value).expanduser()


def venv_python(venv):
    return expanded_path(venv) / "bin" / "python"


def default_venv():
    return os.environ.get("TSRS_KOKORO_VENV") or DEFAULT_VENV


def state_dir():
    return expanded_path(
        os.environ.get("TSRS_KOKORO_STATE_DIR")
        or "~/Library/Application Support/Tri-State Relay Service/kokoro"
    )


def socket_path():
    return state_dir() / "kokoro.sock"


def pid_path():
    return state_dir() / "kokoro.pid"


def log_path():
    return state_dir() / "kokoro.log"


def ensure_state_dir():
    directory = state_dir()
    directory.mkdir(parents=True, exist_ok=True)
    stat = directory.stat()
    if stat.st_uid != os.getuid():
        raise RuntimeError(f"Kokoro state directory is not owned by the current user: {directory}")
    directory.chmod(0o700)
    return directory


def infer_venv_from_python(python):
    path = expanded_path(python)
    if path.name != "python":
        return None
    candidate = path.parent.parent
    if (candidate / "pyvenv.cfg").exists():
        return str(candidate)
    return None


def resolve_python(args):
    if args.python:
        python = args.python
        return python, args.venv or infer_venv_from_python(python)

    if args.venv:
        return str(venv_python(args.venv)), args.venv

    env_python = os.environ.get("TSRS_KOKORO_PYTHON")
    if env_python:
        return env_python, infer_venv_from_python(env_python)

    venv = default_venv()
    candidate = venv_python(venv)
    if candidate.exists():
        return str(candidate), venv

    return "python3", None


def kokoro_environment(venv, enable_mps):
    env = os.environ.copy()
    if venv:
        venv_path = str(expanded_path(venv))
        env["VIRTUAL_ENV"] = venv_path
        env["PATH"] = f"{venv_path}/bin:{env.get('PATH', '')}"
    if enable_mps or os.environ.get("TSRS_KOKORO_ENABLE_MPS") == "1":
        env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
    return env


def install_hint():
    return (
        "Install Kokoro in an isolated venv, then rerun this command:\n"
        "  brew install uv espeak-ng\n"
        f"  uv venv --python 3.11 {DEFAULT_VENV}\n"
        f"  uv pip install --python {DEFAULT_VENV}/bin/python kokoro==0.9.4\n"
        "The first synthesis may download Kokoro-82M weights and the spaCy English model."
    )


def preflight_python(python, env):
    try:
        result = subprocess.run(
            [
                python,
                "-c",
                "import importlib.util; raise SystemExit(0 if importlib.util.find_spec('kokoro') else 1)",
            ],
            check=False,
            capture_output=True,
            text=True,
            env=env,
        )
    except OSError as error:
        return f"could not start Python for Kokoro: {error}"

    if result.returncode == 0:
        return None

    detail = (result.stderr or result.stdout).strip().splitlines()
    suffix = f": {detail[-1]}" if detail else ""
    return f"Kokoro is not installed for {python}{suffix}"


def synthesis_output_path(output_file):
    if output_file.suffix.lower() == ".wav":
        return output_file, None

    handle = tempfile.NamedTemporaryFile(
        prefix=".kokoro-",
        suffix=".wav",
        dir=str(output_file.parent),
        delete=False,
    )
    handle.close()
    return Path(handle.name), Path(handle.name)


def publish_synthesized_audio(synthesis_file, output_file):
    if synthesis_file == output_file:
        return True

    os.replace(synthesis_file, output_file)
    return True


def read_pid():
    try:
        return int(pid_path().read_text(encoding="utf-8").strip())
    except (OSError, ValueError):
        return None


def process_is_running(pid):
    if not pid:
        return False
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def pid_looks_like_kokoro_server(pid):
    try:
        result = subprocess.run(
            ["/bin/ps", "-p", str(pid), "-o", "command="],
            check=False,
            capture_output=True,
            text=True,
        )
    except OSError:
        return False

    command = result.stdout
    return result.returncode == 0 and "kokoro-voice-command" in command and "server" in command and "run" in command


def socket_request(payload, timeout=30.0):
    path = socket_path()
    client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    client.settimeout(timeout)
    try:
        client.connect(str(path))
        client.sendall(json.dumps(payload, separators=(",", ":")).encode("utf-8") + b"\n")
        chunks = []
        while True:
            chunk = client.recv(65536)
            if not chunk:
                break
            chunks.append(chunk)
            if b"\n" in chunk:
                break
    finally:
        client.close()

    data = b"".join(chunks).split(b"\n", 1)[0]
    if not data:
        raise RuntimeError("Kokoro server returned no response")
    response = json.loads(data.decode("utf-8"))
    if not response.get("ok"):
        raise RuntimeError(response.get("error") or "Kokoro server request failed")
    return response


def server_is_responsive():
    try:
        response = socket_request({"action": "status", "version": SERVER_PROTOCOL_VERSION}, timeout=1.0)
    except Exception:
        return False
    return response.get("status") == "ready"


def start_server(args, python, venv, env):
    ensure_state_dir()
    if process_is_running(read_pid()) and server_is_responsive():
        return

    stale_socket = socket_path()
    if stale_socket.exists():
        stale_socket.unlink()

    command = [
        python,
        str(Path(__file__).resolve()),
        "server",
        "run",
        "--idle-timeout",
        str(args.server_idle_timeout),
    ]
    if venv:
        command += ["--venv", venv]
    if args.mps:
        command.append("--mps")

    with open(log_path(), "ab") as log:
        subprocess.Popen(
            command,
            stdin=subprocess.DEVNULL,
            stdout=log,
            stderr=log,
            env=env,
            start_new_session=True,
            close_fds=True,
        )

    deadline = time.time() + float(args.server_start_timeout)
    last_error = None
    while time.time() < deadline:
        try:
            if server_is_responsive():
                return
        except Exception as error:
            last_error = error
        time.sleep(0.1)

    raise RuntimeError(f"Kokoro server did not become ready; see {log_path()}" + (f": {last_error}" if last_error else ""))


def stop_server(timeout=1.0):
    try:
        socket_request({"action": "shutdown", "version": SERVER_PROTOCOL_VERSION}, timeout=1.0)
    except Exception:
        pass

    deadline = time.time() + timeout
    pid = read_pid()
    while pid and process_is_running(pid) and time.time() < deadline:
        time.sleep(0.1)

    if pid and process_is_running(pid) and pid_looks_like_kokoro_server(pid):
        os.kill(pid, signal.SIGTERM)

    for path in [socket_path(), pid_path()]:
        try:
            if path.exists():
                path.unlink()
        except OSError:
            pass


def server_status():
    pid = read_pid()
    responsive = server_is_responsive() if pid and process_is_running(pid) else False
    return {
        "running": bool(pid and process_is_running(pid)),
        "responsive": responsive,
        "pid": pid,
        "socket": str(socket_path()),
        "log": str(log_path()),
    }


def send_server_response(connection, response):
    try:
        connection.sendall(json.dumps(response, separators=(",", ":")).encode("utf-8") + b"\n")
    except BrokenPipeError:
        return


def write_wav_with_kokoro(text, output_file, voice_id, language=None, speed=1.0):
    import wave
    import numpy as np
    from kokoro import KPipeline

    lang = language or voice_id[0]
    pipeline = write_wav_with_kokoro.pipelines.get(lang)
    if pipeline is None:
        pipeline = KPipeline(lang_code=lang, repo_id="hexgrad/Kokoro-82M")
        write_wav_with_kokoro.pipelines[lang] = pipeline

    with wave.open(str(output_file), "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(24000)
        for result in pipeline(text, voice=voice_id, speed=speed, split_pattern=r"\n+"):
            if result.audio is None:
                continue
            wav_file.writeframes((result.audio.numpy() * 32767).astype(np.int16).tobytes())


write_wav_with_kokoro.pipelines = {}


def server_loop(idle_timeout):
    ensure_state_dir()
    path = socket_path()
    if path.exists():
        path.unlink()

    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(str(path))
    path.chmod(0o600)
    server.listen(8)
    server.settimeout(0.5)
    pid_path().write_text(str(os.getpid()), encoding="utf-8")

    shutting_down = False
    last_request_at = time.time()
    try:
        while not shutting_down:
            if idle_timeout > 0 and time.time() - last_request_at > idle_timeout:
                break

            try:
                connection, _ = server.accept()
            except socket.timeout:
                continue

            with connection:
                last_request_at = time.time()
                request = connection.recv(1024 * 1024)
                try:
                    payload = json.loads(request.split(b"\n", 1)[0].decode("utf-8"))
                    action = payload.get("action")
                    if action == "status":
                        response = {"ok": True, "status": "ready", "pid": os.getpid()}
                    elif action == "shutdown":
                        response = {"ok": True, "status": "stopping", "pid": os.getpid()}
                        shutting_down = True
                    elif action == "synthesize":
                        text_file = expanded_path(payload["text_file"])
                        output_file = expanded_path(payload["output_file"])
                        voice_id = payload.get("voice_id") or DEFAULT_VOICE_ID
                        language = payload.get("language")
                        speed = float(payload.get("speed") or 1.0)
                        output_file.parent.mkdir(parents=True, exist_ok=True)
                        synthesis_file, temporary_wav = synthesis_output_path(output_file)
                        write_wav_with_kokoro(text_file.read_text(encoding="utf-8"), synthesis_file, voice_id, language=language, speed=speed)
                        publish_synthesized_audio(synthesis_file, output_file)
                        if temporary_wav and temporary_wav.exists():
                            temporary_wav.unlink()
                        response = {"ok": True, "status": "synthesized", "output_file": str(output_file)}
                    else:
                        response = {"ok": False, "error": f"unsupported Kokoro server action: {action}"}
                except Exception as error:
                    response = {"ok": False, "error": str(error)}
                send_server_response(connection, response)
    finally:
        server.close()
        for cleanup_path in [path, pid_path()]:
            try:
                if cleanup_path.exists():
                    cleanup_path.unlink()
            except OSError:
                pass


def run_server(argv):
    parser = argparse.ArgumentParser(description="Manage the local TSRS Kokoro helper server.")
    subparsers = parser.add_subparsers(dest="command")

    run_parser = subparsers.add_parser("run")
    run_parser.add_argument("--venv", default=None)
    run_parser.add_argument("--mps", action="store_true")
    run_parser.add_argument("--idle-timeout", type=float, default=float(os.environ.get("TSRS_KOKORO_SERVER_IDLE_TIMEOUT", DEFAULT_SERVER_IDLE_TIMEOUT_SECONDS)))

    stop_parser = subparsers.add_parser("stop")
    stop_parser.add_argument("--timeout", type=float, default=1.0)
    subparsers.add_parser("status")

    args = parser.parse_args(argv)
    if args.command == "run":
        server_loop(args.idle_timeout)
        return 0
    if args.command == "stop":
        ensure_state_dir()
        stop_server(timeout=args.timeout)
        print(json.dumps(server_status(), separators=(",", ":")))
        return 0
    if args.command == "status":
        ensure_state_dir()
        print(json.dumps(server_status(), separators=(",", ":")))
        return 0
    parser.print_help()
    return 2


def run_voices(argv):
    parser = argparse.ArgumentParser(description="List Kokoro voice ids for TSRS line-voice assignment.")
    parser.add_argument("--language", choices=sorted({voice[0] for voice in VOICE_IDS}))
    args = parser.parse_args(argv)

    ids = VOICE_IDS
    if args.language:
        ids = [voice for voice in ids if voice.startswith(args.language)]
    print(json.dumps(ids))
    return 0


def run_speech(argv):
    parser = argparse.ArgumentParser(description="Generate a TSRS voice-command WAV file with Kokoro.")
    parser.add_argument("--text-file", required=True)
    parser.add_argument("--output-file", required=True)
    parser.add_argument("--voice-id", default=DEFAULT_VOICE_ID)
    parser.add_argument("--language", choices=sorted({voice[0] for voice in VOICE_IDS}))
    parser.add_argument("--speed", type=float, default=1.0)
    parser.add_argument("--venv", default=None)
    parser.add_argument("--python", default=None)
    parser.add_argument("--mps", action="store_true", help="Enable PyTorch MPS fallback for Apple Silicon.")
    parser.add_argument("--no-server", action="store_true")
    parser.add_argument("--server-start-timeout", type=float, default=float(os.environ.get("TSRS_KOKORO_SERVER_START_TIMEOUT", 30)))
    parser.add_argument("--server-idle-timeout", type=float, default=float(os.environ.get("TSRS_KOKORO_SERVER_IDLE_TIMEOUT", DEFAULT_SERVER_IDLE_TIMEOUT_SECONDS)))
    args = parser.parse_args(argv)

    text_file = expanded_path(args.text_file)
    output_file = expanded_path(args.output_file)
    if not text_file.exists():
        print(f"text file does not exist: {text_file}", file=sys.stderr)
        return 2

    output_file.parent.mkdir(parents=True, exist_ok=True)

    python, venv = resolve_python(args)
    env = kokoro_environment(venv, args.mps)
    error = preflight_python(python, env)
    if error:
        print(error, file=sys.stderr)
        print(install_hint(), file=sys.stderr)
        return 2

    if not args.no_server and os.environ.get("TSRS_KOKORO_DISABLE_SERVER") != "1":
        try:
            start_server(args, python, venv, env)
            socket_request(
                {
                    "action": "synthesize",
                    "version": SERVER_PROTOCOL_VERSION,
                    "text_file": str(text_file),
                    "output_file": str(output_file),
                    "voice_id": args.voice_id,
                    "language": args.language,
                    "speed": args.speed,
                },
                timeout=max(30.0, args.server_start_timeout),
            )
            return 0
        except Exception as error:
            print(f"Kokoro server unavailable; falling back to one-shot synthesis: {error}", file=sys.stderr)

    synthesis_file, temporary_wav = synthesis_output_path(output_file)
    command = [
        python,
        "-m",
        "kokoro",
        "--input-file",
        str(text_file),
        "--output-file",
        str(synthesis_file),
        "--voice",
        args.voice_id,
        "--speed",
        str(args.speed),
    ]
    if args.language:
        command += ["--language", args.language]

    try:
        result = subprocess.run(command, check=False, env=env)
    except OSError as error:
        print(f"Kokoro could not start: {error}", file=sys.stderr)
        return 1

    try:
        if result.returncode != 0:
            return result.returncode
        if not synthesis_file.exists() or synthesis_file.stat().st_size == 0:
            print(f"Kokoro did not write audio output: {synthesis_file}", file=sys.stderr)
            return 1
        if not publish_synthesized_audio(synthesis_file, output_file):
            return 1
        if not output_file.exists() or output_file.stat().st_size == 0:
            print(f"Kokoro did not write audio output: {output_file}", file=sys.stderr)
            return 1
        return 0
    finally:
        if temporary_wav and temporary_wav.exists():
            temporary_wav.unlink()


def run_self_test():
    class KokoroWrapperTests(unittest.TestCase):
        def test_voices_are_json_and_filterable(self):
            stdout = io.StringIO()
            with contextlib.redirect_stdout(stdout):
                self.assertEqual(run_voices(["--language", "a"]), 0)
            ids = json.loads(stdout.getvalue())
            self.assertIn(DEFAULT_VOICE_ID, ids)
            self.assertTrue(all(voice.startswith("a") for voice in ids))

        def test_missing_python_fails_with_install_hint(self):
            with tempfile.TemporaryDirectory() as directory:
                text = Path(directory) / "relay.txt"
                output = Path(directory) / "relay.wav"
                text.write_text("hello", encoding="utf-8")
                stderr = io.StringIO()
                with contextlib.redirect_stderr(stderr):
                    result = run_speech([
                        "--text-file",
                        str(text),
                        "--output-file",
                        str(output),
                        "--python",
                        str(Path(directory) / "missing-python"),
                    ])
            self.assertEqual(result, 2)
            self.assertIn("Install Kokoro in an isolated venv", stderr.getvalue())

        def test_venv_resolution_sets_environment_path(self):
            env = kokoro_environment("~/example-venv", True)
            self.assertEqual(env["VIRTUAL_ENV"], str(expanded_path("~/example-venv")))
            self.assertTrue(env["PATH"].startswith(f"{expanded_path('~/example-venv')}/bin:"))
            self.assertEqual(env["PYTORCH_ENABLE_MPS_FALLBACK"], "1")

        def test_non_wav_output_uses_temporary_wav(self):
            with tempfile.TemporaryDirectory() as directory:
                output = Path(directory) / "relay.audio"
                synthesis_file, temporary_wav = synthesis_output_path(output)
                self.assertNotEqual(synthesis_file, output)
                self.assertEqual(synthesis_file.suffix, ".wav")
                self.assertEqual(temporary_wav, synthesis_file)

        def test_wav_output_does_not_need_conversion(self):
            with tempfile.TemporaryDirectory() as directory:
                output = Path(directory) / "relay.wav"
                synthesis_file, temporary_wav = synthesis_output_path(output)
                self.assertEqual(synthesis_file, output)
                self.assertIsNone(temporary_wav)

        def test_publish_synthesized_audio_renames_temporary_wav(self):
            with tempfile.TemporaryDirectory() as directory:
                synthesis = Path(directory) / ".kokoro-test.wav"
                output = Path(directory) / "relay.audio"
                synthesis.write_bytes(b"RIFF-test")

                self.assertTrue(publish_synthesized_audio(synthesis, output))
                self.assertFalse(synthesis.exists())
                self.assertEqual(output.read_bytes(), b"RIFF-test")

        def test_install_hint_has_complete_setup_path(self):
            hint = install_hint()
            self.assertIn("brew install uv espeak-ng", hint)
            self.assertIn("uv venv --python 3.11", hint)
            self.assertIn("kokoro==0.9.4", hint)

        def test_broken_pipe_does_not_crash_server_response(self):
            class BrokenConnection:
                def sendall(self, data):
                    raise BrokenPipeError()

            send_server_response(BrokenConnection(), {"ok": True})

    suite = unittest.defaultTestLoader.loadTestsFromTestCase(KokoroWrapperTests)
    return 0 if unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful() else 1


def main():
    if len(sys.argv) > 1 and sys.argv[1] == "--self-test":
        return run_self_test()
    if len(sys.argv) > 1 and sys.argv[1] == "server":
        return run_server(sys.argv[2:])
    if len(sys.argv) > 1 and sys.argv[1] == "voices":
        return run_voices(sys.argv[2:])
    return run_speech(sys.argv[1:])


if __name__ == "__main__":
    raise SystemExit(main())
